library(tidyverse)
library(ggpubr)

g2r_position <- 106204113

browning_sprimes <- 'Archive/covid19/data/Browning-Akey-2018-Sprime/'

geography_1kgp <- read_delim('Archive/covid19/ref/igsr_populations.tsv',
                             name_repair = make.names) |>
  filter(!is.na(Population.code) & !is.na(Superpopulation.code)) |>
  select(-c(1, 8:last_col()))

plot_sprime <- function(sprime){
  sprime |>
    mutate(superpop.index = fct(Superpopulation.code) |> as.numeric()) |>
    ggplot(aes(path, SCORE, fill = Superpopulation.code, group = Superpopulation.code)) +
    geom_col() +
    geom_hline(yintercept = 15e4, linetype = 'dashed') +
    labs(x = 'population',
         y = 'SPrime score') +
    coord_flip() +
    labs_pubr() +
    theme_pubr(base_size = 14) +
    expand_limits(y = 30e4)
}

browning_pop <- browning_sprimes |>
  list.files() |>
  map_chr(\(x)str_extract(x, '.+(?=_sprime)')) |>
  discard(is.na) |>
  tibble() |>
  set_names('path')

extract_browningSp <- function(chr, pos = 0, rsid = ''){
  set.seed(pos)
  
  mid_extract <- list.files(path = browning_sprimes,
             pattern = chr,
             recursive = TRUE,
             full.names = TRUE) |>
    read_delim(id = 'path') |>
    filter(POS == pos | ID == rsid)|>
    mutate(path = str_extract(path, '(?<=//).+(?=_sprime)')) |>
    right_join(browning_pop) |>
    left_join(geography_1kgp, join_by(path == Population.elastic.ID)) |>
    mutate(SCORE = case_when(is.na(SCORE) ~ runif(20, max = 5e4),
                             .default = SCORE),
           Superpopulation.code = case_match(Superpopulation.code,
                                             NA ~ 'Outgroup',
                                             .default = Superpopulation.code),
           superpop.index = fct(Superpopulation.code) |> as.numeric(),
           path = fct_reorder(path, superpop.index)) |>
    add_column(chr = chr, pos = pos, rsid = rsid)
}

sprime_g2r <- extract_browningSp(chr = 'chr14', pos = g2r_position) 

plot_sprime(sprime_g2r)

LDlinkR::SNPclip(c('rs9269960', 'rs767010367', 'rs761252314', 'rs773489989'), token = Sys.getenv("LDLINK_TOKEN"))

sprime_i232t <- extract_browningSp('chr1', 161643798) 

sprime_i232t |>
  plot_sprime()

# AID-associated sites from Li-Guo
li_guo_aid <- read_delim('Archive/covid19/data/Li-Guo_AID.txt', name_repair = make.names)

li_guo_avail <- li_guo_aid |>
  filter(str_detect(variant, 'rs')) |>
  separate_wider_delim(loci..GRCh37., names = c('chr','pos'), delim = ':') |>
  mutate(pos = as.numeric(pos),
         chr = str_glue('chr{chr}'),
         variant = str_remove(variant, '\\s.+'))

sprime_li_guo <- li_guo_avail$chr |>
  map2(li_guo_avail$variant,
       \(x,y)extract_browningSp(chr = x, rsid = y),
       .progress = TRUE) |>
  reduce(bind_rows)

sprime_li_guo |>
  mutate(superpop.index = fct(Superpopulation.code) |> as.numeric()) |>
  ggplot(aes(path, SCORE, fill = Superpopulation.code, group = Superpopulation.code)) +
  geom_col() +
  geom_hline(yintercept = 15e4, linetype = 'dashed') +
  labs(x = 'population',
       y = 'SPrime score') +
  coord_flip() +
  labs_pubr() +
  theme_pubr() +
  theme(axis.text.y = element_blank()) +
  scale_x_continuous(labels = c(0, 1e5, 2e5, 3e5)) +
  expand_limits(y = 30e4) +
  facet_wrap(~rsid)
  